#! /usr/bin/env python
# coding=utf-8



from barcode_checking import *
from Bio import SeqIO
import os
import shutil
import sys
import argparse
from argparse import RawTextHelpFormatter
import threading


parser = argparse.ArgumentParser(
    description='''
    使用已知数据库中的基因条码序列对 NGSpeciesID的cluster 进行修正
    这一步会比对已知数据库中 最近的10条序列到目标cluster序列 然后通过多序列比对 找到
    用法:
    barcode_fix_seq.py -i ngsp_filter_out -t 2 -c 4 -b ../00_work_simple_nt/simple_nt -r ../00_work_simple_nt/simple_nt.fa
    由大天才于2022年3月24日创建于浙江农业大学''',formatter_class=RawTextHelpFormatter)



parser.add_argument('-i',
                help='必须给定，NGSpeciesID标本的路径')

parser.add_argument('-t',
                help='可选, 运行一条序列进行校正的核数 默认为2')

parser.add_argument('-c',
                help='可选, 同时运行多少条序列 默认为4')

parser.add_argument('-b',
                help='必须给定，参考序列的blast库')
parser.add_argument('-r',
                help='必须给定，参考序列的fasta文件')



args = parser.parse_args()

if not args.i or not args.b or not args.r:

    parser.print_help()

    sys.exit()

infile = args.i

search_path = args.b

search_fasta = args.r

if not args.t:
	num_threads = 2
else:
	num_threads = int(args.t)

if not args.c:
	num_cores = 4
else:
	num_cores = int(args.c)
	

# 对一条序列运行 megablast 之后使用mafft 进行比对 使用比对的结果 对原序列进行修正
# infile = r'ngsp_out_test/COI_F_1A_R_2H/consensus_reference_2.fasta'

# search_path = '../00_work_simple_nt/simple_nt'

# search_fasta = '../00_work_simple_nt/simple_nt.fa'




# 读取 所有参考序列
print('读取所有参考序列')
tmp_seq_dic = {}
for i in SeqIO.parse(search_fasta,'fasta'):
	tmp_seq_dic[i.name]= str(i.seq)

print('读取完毕')

	

# 多线程运行


index = 0


thread_pool = []


def fix_one_sequence(target_file_path):

	global index
	global core
	#print(target_file_path)
	index = index + 1

	tmp_file_name = '.'.join(target_file_path.split('.')[:-1])
	if not os.path.exists(tmp_file_name):

		try:
			os.mkdir(tmp_file_name)
		except:
			pass

	#if translation!=False:
		# 如果 是可翻译的 检查是否该蛋白与已知数据库中的首尾匹配 如果匹配 才会认为是 good quality的
	os.system('blastn -query %s  -db %s  -out %s/blastn_res -num_threads %s  -evalue 1e-5 -outfmt 6 -max_target_seqs 10'  # -task dc-megablast
		% (target_file_path, search_path, tmp_file_name, num_threads))

	# 建立用于序列比对的文件
	fix_file = open(target_file_path+'.fixed','w')
	target_seq =[]
	with open(tmp_file_name+'/blastn_res') as fila:	
		for i in fila:
			k = i.strip().split('\t')
			if len(k)>1:
				target_seq.append([k[1],tmp_seq_dic[k[1]]])

	# 如果只有一个target seq 也不会做 必须高于1个
	seq = next(SeqIO.parse(target_file_path,'fasta'))
	if len(target_seq)>=1:
		with open(tmp_file_name+'/multi_ali.fasta','w') as fila:
			fila.write('>'+seq.name+'\n'+str(seq.seq)+'\n')
			for i in target_seq:
				fila.write('>'+i[0]+'\n'+i[1]+'\n')

		# 运行mafft
		#print('运行mafft')
		os.system('mafft --localpair --thread %s  %s/multi_ali.fasta > %s/mafft_out.fasta ' % (num_threads,tmp_file_name,tmp_file_name))

		#print('根据 mafft 结果修正序列')


		fix_seq = mafft2GapSeq(tmp_file_name+'/mafft_out.fasta')
		
		#print(fix_seq)
		if fix_seq!=False:
			fix_file.write(fix_seq[0]+'\t'+';'.join(fix_seq[1]))
		else:
			fix_file.write(''+'\tunaligned')
	else:

		fix_file.write(''+'\tunmatch')

	fix_file.close()



	if index >= len(thread_pool):
		return None
	else:
		thread_pool[index].start()

	return	None

n = 0

for i in os.listdir(infile):
	p1 = infile+'/'+i
	for j in os.listdir(p1):
		if j.split('.')[-1] == 'fasta':
			inpath = p1+'/'+j
			th = threading.Thread(target=fix_one_sequence, args=[inpath], name='th_'+str(n))
			thread_pool.append(th)



for i in range(num_cores):
	if len(thread_pool)>i:
		thread_pool[i].start()
